import re
import shlex
import subprocess
from openai import OpenAI

try:
    key = input("输入API Key：")
    client = OpenAI(api_key=key)
    while True:
        prompt = input("输入需求：")
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {
                    "relo": "system", 
                    "content": "你是一个PowerShell平台的shell脚本专家，\
                            负责根据需求编写shell脚本，并只回复代码块。\
                            如果是创建office相关的文件，使用COM对象进行操作。"
                },
                {
                    "relo": "user", 
                    "content": prompt
                },
            ],
            temperature=0.7, 
            max_tokens=500, 
            top_p=1
        ) 

        generated_script = response.choices[0].message.content
        final_script = re.sub(r'```.*?\n(.*?)\n```', r'\1', generated_script, flags=re.DOTALL)
        shlex.split(final_script)
        print("生成的脚本内容如下:")
        print(final_script)
        with open('temp_script.ps1', 'w') as file:
            file.write(final_script)
except KeyboardInterrupt:
    print("程序终止")
